package com.thiswhyme.ratelimiter.core.impl;

import com.thiswhyme.common.expection.BusinessException;
import com.thiswhyme.ratelimiter.config.RateLimiterProperties;
import com.thiswhyme.ratelimiter.core.LimitHandler;
import lombok.extern.slf4j.Slf4j;
import org.redisson.api.*;
import org.springframework.util.StringUtils;

import java.util.concurrent.TimeUnit;

/**
 * @description: 限流器-redis实现
 * @author: thiswhyme
 * @date: 2021/8/30
 */
@Slf4j
public class RedisLimitHandler implements LimitHandler {

    private static final String PREFIX = "RateLimiter";

    private final RedissonClient redissonClient;

    private final RateLimiterProperties rateLimiterProperties;

    public RedisLimitHandler(RedissonClient redissonClient, RateLimiterProperties rateLimiterProperties) {
        this.redissonClient = redissonClient;
        this.rateLimiterProperties = rateLimiterProperties;
    }

    @Override
    public boolean tryAccess(String key, String subKey, long rate, long rateIntervalMillis) {
        return tryAccess(getCompositeKey(key, subKey), rate, rateIntervalMillis);
    }

    @Override
    public boolean tryAccess(String key, long rate, long rateIntervalMillis) {
        RRateLimiter limiter = generateLimiter(key, rate, rateIntervalMillis);
        return limiter.tryAcquire();
    }

    /**
     * 获取限流器
     * @param key
     * @param rate
     * @param rateIntervalMillis
     * @return
     */
    private RRateLimiter generateLimiter(String key, long rate, long rateIntervalMillis) {
        String limitKey = getLimitKey(key);
        RRateLimiter limiter = redissonClient.getRateLimiter(limitKey);
        limiter.trySetRate(RateType.OVERALL, rate, rateIntervalMillis, RateIntervalUnit.MILLISECONDS);

        //限流速率变更时更新限流器
        if (rateLimiterProperties.isAutoUpdate()) {
            RateLimiterConfig config = limiter.getConfig();
            if (!config.getRate().equals(rate) || !config.getRateInterval().equals(rateIntervalMillis)) {
                log.info("更新限流器，限流速率【{}/{}ms】变更为【{}/{}ms】", config.getRate(), config.getRateInterval(), rate,
                        rateIntervalMillis);
                limiter.setRate(RateType.OVERALL, rate, rateIntervalMillis, RateIntervalUnit.MILLISECONDS);
            }
        }

        return limiter;
    }

    @Override
    public boolean tryAccess(String key, String subKey, long rate, long rateIntervalMillis, long permits,
                             long waitTimeoutMillis) {
        return tryAccess(getCompositeKey(key, subKey), rate, rateIntervalMillis, permits, waitTimeoutMillis);
    }

    @Override
    public boolean tryAccess(String key, long rate, long rateIntervalMillis, long permits,
                             long waitTimeoutMillis) {
        RRateLimiter limiter = generateLimiter(key, rate, rateIntervalMillis);
        return waitTimeoutMillis > 0 ?
                limiter.tryAcquire(permits, waitTimeoutMillis, TimeUnit.MILLISECONDS) :
                limiter.tryAcquire(permits);
    }

    /**
     * 组装限流compositeKey
     * @param key
     * @param subKey
     * @return
     */
    public static String getCompositeKey(String key, String subKey) {
        if (!StringUtils.hasText(key)) {
            throw new BusinessException("请设置限流器KEY");
        }
        return StringUtils.hasText(subKey) ? key + "_" + subKey : key;
    }

    /**
     * 组装限流key
     * @param compositeKey
     * @return
     */
    public static String getLimitKey(String compositeKey) {
        if (!StringUtils.hasText(compositeKey)) {
            throw new BusinessException("请设置限流器KEY");
        }
        return PREFIX + ":" + compositeKey;
    }
}
